"""
A function that plots data and allows the user to click to add points and select points
to remove them.  The function returns the list of points.
"""

from matplotlib.widgets import RectangleSelector
from pylab import subplot, arange, plot, sin, cos, pi, show
import sys

#----------------------------------------------------------------
#-- CLASSES AND FUNCTIONS --
#----------------------------------------------------------------
def rectangle_select_callback(event1, event2):
    'event1 and event2 are the press and release events'
    x1, y1 = event1.xdata, event1.ydata
    x2, y2 = event2.xdata, event2.ydata
    #print "(%3.2f, %3.2f) --> (%3.2f, %3.2f)"%(x1,y1,x2,y2)
    #print " The button you used were: ",event1.button, event2.button
    add_or_remove_pt(x1,y1,x2,y2)

def add_or_remove_pt(x1,y1,x2,y2):
    '''
    Maintain a list of points clicked on but not selected out.
    '''
    # find out if any points are within the rectangle defined by the arguements.
    ux = max(x1,x2)
    lx = min(x1,x2)
    uy = max(y1,y2)
    ly = min(y1,y2)

    remove_pts = []
    for x,y in add_or_remove_pt.list_of_pts:
        if x >= lx and x <= ux and\
           y >= ly and y <= uy:
            remove_pts.append((x,y))
    for pt in remove_pts:
        add_or_remove_pt.list_of_pts.remove(pt)

    # if we didn't remove any points, lets add one
    if len(remove_pts) == 0:
        add_or_remove_pt.list_of_pts.append( ((ux+lx)/2.0, (uy+ly)/2.0) )
    xlist = []
    ylist = []
    for x,y in add_or_remove_pt.list_of_pts:
        xlist.append(x)
        ylist.append(y)
    add_or_remove_pt.line1.set_data(xlist,ylist)
    show()
        

def point_maker(datax,datay):
    ax=subplot(111)                    # make a new plotingrange

    if len(datax) > 2:
        add_or_remove_pt.line1, = plot(datax[0:2],datay[0:2],'o')   # plot something this will be overridden by clicking
    else:
        raise RuntimeError("not enough data points in xdata/ydata.")
    plot(datax,datay)

    add_or_remove_pt.list_of_pts = []
    # drawtype is 'box' or 'line' or 'none'
    LS = RectangleSelector(ax, rectangle_select_callback,
                      drawtype='box',useblit=True)
    show()
    return add_or_remove_pt.list_of_pts
